"""
ChromaDB 配置和操作类
用于管理向量数据库的连接、初始化和CRUD操作
"""

import os
from typing import List, Dict, Any, Optional, Union
from uuid import uuid4

import chromadb
from chromadb.config import Settings as ChromaSettings
from chromadb.utils import embedding_functions

from config.logging_config import get_logger
from constant.settings import settings

# 使用全局logger设置函数
logger = get_logger(__name__)


class ChromaDBConfig:
    """ChromaDB 配置和操作类"""

    def __init__(self):
        """初始化ChromaDB配置"""
        self.client = None
        self.collection = None
        self.embedding_function = None
        self._initialize()

    def _initialize(self):
        """初始化ChromaDB客户端和集合"""
        try:
            # 设置ChromaDB数据存储路径
            chroma_db_path = settings.CHROMA_DB_PATH
            os.makedirs(chroma_db_path, exist_ok=True)

            # 初始化ChromaDB客户端
            self.client = chromadb.PersistentClient(
                path=chroma_db_path,
                settings=ChromaSettings(
                    anonymized_telemetry=False,
                    allow_reset=True
                )
            )

            # 设置嵌入函数 (使用EmbFactory创建)
            try:
                # 延迟导入避免循环依赖
                from emb.EmbFactory import EmbFactory

                # 构建模型代码
                model_code = settings.CHROMA_EMBEDDING_MODEL

                # 使用工厂方法创建嵌入函数
                self.embedding_function = EmbFactory.create_embedding_function(
                    emb_code=model_code,
                )

                if self.embedding_function:
                    logger.info(f"成功通过EmbFactory创建嵌入函数，模型: {model_code}")
                else:
                    raise Exception("EmbFactory返回None")

            except Exception as e:
                logger.warning(f"通过EmbFactory创建嵌入函数失败: {e}")
                logger.info("回退到默认嵌入函数")
                # 回退到默认嵌入函数
                self.embedding_function = embedding_functions.DefaultEmbeddingFunction()

            # 获取或创建默认集合
            collection_name = settings.CHROMA_COLLECTION_NAME

            # 首先检查集合是否存在
            existing_collections = [col.name for col in self.client.list_collections()]
            collection_exists = collection_name in existing_collections
            if collection_exists:
                self.collection = self.client.get_collection(
                    name=collection_name,
                    embedding_function=self.embedding_function
                )
                logger.info(f"成功连接到现有集合: {collection_name}")
            else:
                # 如果集合不存在，则创建新集合
                self.collection = self.client.create_collection(
                    name=collection_name,
                    embedding_function=self.embedding_function,
                    metadata={"description": "AI Web 向量数据存储集合"}
                )
                logger.info(f"成功创建新集合: {collection_name}")

        except Exception as e:
            logger.error(f"ChromaDB 初始化失败: {e}")
            raise

    def add_documents(self,
                      documents: List[str],
                      metadatas: List[Dict[str, Any]],
                      ids: Optional[List[str]] = None,
                      scene_id: Optional[str] = None) -> bool:
        """
        添加文档到向量数据库
        
        Args:
            documents: 文档内容列表
            metadatas: 文档元数据列表
            ids: 文档ID列表（可选，如果不提供会自动生成）
            scene_id: 场景ID（用于索引优化）
            
        Returns:
            bool: 添加是否成功
        """
        try:
            if not documents or not metadatas:
                raise ValueError("文档内容和元数据不能为空")

            if len(documents) != len(metadatas):
                raise ValueError("文档内容和元数据数量必须一致")

            # 如果没有提供ID，自动生成
            if ids is None:
                ids = [str(uuid4()) for _ in documents]
            elif len(ids) != len(documents):
                raise ValueError("文档ID和文档内容数量必须一致")

            # 为所有元数据添加scene_id
            if scene_id:
                for metadata in metadatas:
                    metadata["scene_id"] = scene_id

            # 添加文档到集合
            self.collection.add(
                documents=documents,
                metadatas=metadatas,
                ids=ids
            )

            logger.info(f"成功添加 {len(documents)} 个文档到ChromaDB")
            return True

        except Exception as e:
            logger.error(f"添加文档失败: {e}")
            return False

    def query_documents(self,
                        query_texts: Union[str, List[str]],
                        scene_id: Optional[str] = None,
                        n_results: int = 10,
                        where: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        查询相似文档
        
        Args:
            query_texts: 查询文本（可以是单个字符串或字符串列表）
            scene_id: 场景ID（用于索引过滤）
            n_results: 返回结果数量
            where: 额外的过滤条件
            
        Returns:
            Dict: 查询结果
        """
        try:
            # 处理查询文本格式
            if isinstance(query_texts, str):
                query_texts = [query_texts]

            # 构建过滤条件
            filter_conditions = where or {}

            # 添加scene_id过滤条件
            if scene_id:
                filter_conditions["scene_id"] = scene_id
                logger.info(f"使用scene_id过滤: {scene_id}")

            # 执行查询
            results = self.collection.query(
                query_texts=query_texts,
                n_results=n_results,
                where=filter_conditions if filter_conditions else None,
                include=["documents", "metadatas", "distances"]
            )

            logger.info(
                f"查询完成，返回 {len(results.get('documents', [])[0] if results.get('documents') else 0)} 个结果")
            return results

        except Exception as e:
            logger.error(f"查询文档失败: {e}")
            return {"documents": [], "metadatas": [], "distances": [], "ids": []}

    def update_documents(self,
                         ids: List[str],
                         documents: Optional[List[str]] = None,
                         metadatas: Optional[List[Dict[str, Any]]] = None,
                         scene_id: Optional[str] = None) -> bool:
        """
        更新文档
        
        Args:
            ids: 要更新的文档ID列表
            documents: 新的文档内容列表（可选）
            metadatas: 新的元数据列表（可选）
            scene_id: 场景ID（用于更新元数据）
            
        Returns:
            bool: 更新是否成功
        """
        try:
            if not ids:
                raise ValueError("文档ID列表不能为空")

            update_params = {"ids": ids}

            if documents is not None:
                if len(documents) != len(ids):
                    raise ValueError("文档内容和ID数量必须一致")
                update_params["documents"] = documents

            if metadatas is not None:
                if len(metadatas) != len(ids):
                    raise ValueError("元数据和ID数量必须一致")

                # 为所有元数据添加scene_id
                if scene_id:
                    for metadata in metadatas:
                        metadata["scene_id"] = scene_id

                update_params["metadatas"] = metadatas

            # 执行更新
            self.collection.update(**update_params)

            logger.info(f"成功更新 {len(ids)} 个文档")
            return True

        except Exception as e:
            logger.error(f"更新文档失败: {e}")
            return False

    def delete_documents(self,
                         ids: Optional[List[str]] = None,
                         where: Optional[Dict[str, Any]] = None,
                         scene_id: Optional[str] = None) -> bool:
        """
        删除文档
        
        Args:
            ids: 要删除的文档ID列表（可选）
            where: 删除条件（可选）
            scene_id: 场景ID（用于批量删除某个场景的所有文档）
            
        Returns:
            bool: 删除是否成功
        """
        try:
            delete_params = {}

            if ids:
                delete_params["ids"] = ids

            # 构建删除条件
            filter_conditions = where or {}
            if scene_id:
                filter_conditions["scene_id"] = scene_id

            if filter_conditions:
                delete_params["where"] = filter_conditions

            if not delete_params:
                raise ValueError("必须指定删除条件（ids 或 where 条件）")

            # 执行删除
            self.collection.delete(**delete_params)

            logger.info("成功删除文档")
            return True

        except Exception as e:
            logger.error(f"删除文档失败: {e}")
            return False

    def get_collection_info(self) -> Dict[str, Any]:
        """
        获取集合信息
        
        Returns:
            Dict: 集合信息
        """
        try:
            count = self.collection.count()
            return {
                "name": self.collection.name,
                "count": count,
                "metadata": self.collection.metadata
            }
        except Exception as e:
            logger.error(f"获取集合信息失败: {e}")
            return {}

    def reset_collection(self) -> bool:
        """
        重置集合（删除所有数据）
        
        Returns:
            bool: 重置是否成功
        """
        try:
            # 删除现有集合
            self.client.delete_collection(name=self.collection.name)

            # 重新创建集合
            collection_name = settings.CHROMA_COLLECTION_NAME
            self.collection = self.client.create_collection(
                name=collection_name,
                embedding_function=self.embedding_function,
                metadata={"description": "AI Web 向量数据存储集合"}
            )

            logger.info("成功重置集合")
            return True

        except Exception as e:
            logger.error(f"重置集合失败: {e}")
            return False

    def search_by_scene(self,
                        scene_id: str,
                        query_text: Optional[str] = None,
                        limit: int = 100) -> Dict[str, Any]:
        """
        基于场景ID进行搜索（先索引后向量搜索的优化方法）
        
        Args:
            scene_id: 场景ID
            query_text: 查询文本（可选，如果提供则进行向量搜索）
            limit: 返回结果数量限制
            
        Returns:
            Dict: 搜索结果
        """
        try:
            if query_text:
                # 场景过滤 + 向量搜索
                return self.query_documents(
                    query_texts=query_text,
                    scene_id=scene_id,
                    n_results=limit
                )
            else:
                # 仅场景过滤
                results = self.collection.get(
                    where={"scene_id": scene_id},
                    limit=limit,
                    include=["documents", "metadatas"]
                )
                logger.info(f"基于scene_id查找到 {len(results.get('documents', []))} 个文档")
                return results

        except Exception as e:
            logger.error(f"基于场景搜索失败: {e}")
            return {"documents": [], "metadatas": [], "ids": []}


# 创建全局ChromaDB配置实例
try:
    chroma_db = ChromaDBConfig()
    logger.info("ChromaDB 全局实例初始化成功")
except Exception as e:
    logger.error(f"ChromaDB 全局实例初始化失败: {e}")
    chroma_db = None
